package com.we.mexico.flink.loan.adm;

import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.serialization.SerializationSchema;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.connector.kafka.sink.KafkaRecordSerializationSchema;
import org.apache.flink.connector.kafka.sink.KafkaSink;
import org.apache.flink.connector.kafka.sink.TopicSelector;
import org.apache.flink.connector.kafka.source.KafkaSource;
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.streaming.connectors.kafka.partitioner.FlinkKafkaPartitioner;
import org.apache.flink.streaming.util.serialization.KeyedSerializationSchema;
import org.apache.flink.util.Collector;
import org.apache.flink.util.OutputTag;

import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.we.mexico.flink.utils.FlinkDataStreamUtils;
import com.we.mexico.flink.utils.WeKafkaPropertyReader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;
import java.util.Properties;

/** ADM 层 墨西哥风控用户借款还款需求; */
public class AdmUserLoanLabel {
    public static Logger LOG = LoggerFactory.getLogger(AdmUserLoanLabel.class);
    public static final String RELEASEPROP = "loanlabel/adm_loan_label.properties";
    public static final String OUTKEYBY = "user_key";

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        WeKafkaPropertyReader paramReader = WeKafkaPropertyReader.init(RELEASEPROP);
        /** Consumer Kafka */
        String srcKafkaTopic = paramReader.getTmpKfkTopic();
        String srcKafkaBootStrapServer = paramReader.getTmpKfkBootStrapServer();
        String srckfkGrupId = paramReader.getTmpKfkGroupId();

        Properties consumProp = new Properties();
        consumProp.setProperty("group.id", srckfkGrupId);
        consumProp.setProperty(
                "key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
        consumProp.setProperty(
                "value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");

        KafkaSource<String> kfkSource =
                KafkaSource.<String>builder()
                        .setBootstrapServers(srcKafkaBootStrapServer)
                        .setTopics(srcKafkaTopic)
                        .setStartingOffsets(OffsetsInitializer.earliest())
                        .setProperties(consumProp)
                        .setValueOnlyDeserializer(new SimpleStringSchema())
                        .build();

        SingleOutputStreamOperator<String> dataStream =
                env.fromSource(kfkSource, WatermarkStrategy.noWatermarks(), "user_loan_and_borrow")
                        .uid("user_loan_and_borrow");

        OutputTag<String> userLoanBorrowOutputTag =
                new OutputTag<>("user_loan_borrow", TypeInformation.of(String.class));
        OutputTag<String> repayPlanOutputTag =
                new OutputTag<>("repay_plan", TypeInformation.of(String.class));

        final SingleOutputStreamOperator<String> process =
                dataStream.process(
                        new ProcessFunction<String, String>() {
                            @Override
                            public void processElement(
                                    String value,
                                    ProcessFunction<String, String>.Context ctx,
                                    Collector<String> out)
                                    throws Exception {
                                final JSONObject jsonObject = JSONObject.parseObject(value);
                                final String streamType = jsonObject.getString("stream_type");
                                final String idNumber = jsonObject.getString("id_number");
                                // 借款和还款公用的 id_number
                                if ("common".equals(streamType) && idNumber != null) {
                                    ctx.output(userLoanBorrowOutputTag, value);
                                    ctx.output(repayPlanOutputTag, value);
                                } else if ("borrow".equals(streamType)) {
                                    ctx.output(userLoanBorrowOutputTag, value);
                                } else if ("repay".equals(streamType)) {
                                    ctx.output(repayPlanOutputTag, value);
                                }
                            }
                        });

        DataStream<String> processSideOutput = process.getSideOutput(userLoanBorrowOutputTag);
        DataStream<String> processSideOutput1 = process.getSideOutput(repayPlanOutputTag);

        final SingleOutputStreamOperator<String> process1 =
                FlinkDataStreamUtils.stringToJsonObject(processSideOutput)
                        .keyBy(data -> data.getString("user_key"))
                        .process(
                                new KeyedProcessFunction<String, JSONObject, String>() {
                                    transient ValueState<JSONObject> borrowInfoValueState;

                                    @Override
                                    public void open(Configuration parameters) throws Exception {
                                        borrowInfoValueState =
                                                getRuntimeContext()
                                                        .getState(
                                                                new ValueStateDescriptor<
                                                                        JSONObject>(
                                                                        "BorrowInfo",
                                                                        JSONObject.class));
                                    }

                                    @Override
                                    public void processElement(
                                            JSONObject value,
                                            KeyedProcessFunction<String, JSONObject, String>.Context
                                                    ctx,
                                            Collector<String> out)
                                            throws Exception {
                                        try {

                                            final JSONObject lastBorrowInfo =
                                                    borrowInfoValueState.value();
                                            final JSONObject currState =
                                                    lastBorrowInfo == null
                                                            ? initState(
                                                                    "user_key",
                                                                    "loan_key",
                                                                    "borrow_money",
                                                                    "borrow_day_num",
                                                                    "period_nos",
                                                                    "borrow_time",
                                                                    "loan_purpose",
                                                                    "created_at",
                                                                    "id_number",
                                                                    "audit_count")
                                                            : lastBorrowInfo;

                                            for (Map.Entry<String, Object> entry :
                                                    value.entrySet()) {
                                                if (entry.getValue() != null
                                                        && !entry.getValue()
                                                                .equals(
                                                                        currState.get(
                                                                                entry.getKey()))) {
                                                    currState.put(entry.getKey(), entry.getValue());
                                                }
                                            }

                                            borrowInfoValueState.update(currState);

                                            if (currState.getString("user_key") != null
                                                    && currState.getString("id_number") != null) {
                                                currState.put(
                                                        "topic_name", "adm_hhl01_kafka_user_loan");
                                                currState.remove("stream_type");
                                                out.collect(
                                                        currState.toString(
                                                                SerializerFeature
                                                                        .WriteMapNullValue));
                                            }
                                        } catch (Exception e) {
                                            LOG.error("BorrowInfo 解析失败了" + e + "原始数据为: " + value);
                                        }
                                    }
                                });

        // user_loan
        final SingleOutputStreamOperator<String> process2 =
                FlinkDataStreamUtils.stringToJsonObject(processSideOutput1)
                        .keyBy(data -> data.getString("user_key"))
                        .process(
                                new KeyedProcessFunction<String, JSONObject, String>() {
                                    transient ValueState<String> idNumberState;
                                    transient MapState<String, JSONObject> loanKeyMapState;

                                    @Override
                                    public void open(Configuration parameters) throws Exception {
                                        idNumberState =
                                                getRuntimeContext()
                                                        .getState(
                                                                new ValueStateDescriptor<String>(
                                                                        "idNumberState",
                                                                        String.class));
                                        loanKeyMapState =
                                                getRuntimeContext()
                                                        .getMapState(
                                                                new MapStateDescriptor<
                                                                        String, JSONObject>(
                                                                        "loanKeyMapState",
                                                                        String.class,
                                                                        JSONObject.class));
                                    }

                                    @Override
                                    public void processElement(
                                            JSONObject value,
                                            KeyedProcessFunction<String, JSONObject, String>.Context
                                                    ctx,
                                            Collector<String> out)
                                            throws Exception {
                                        try {

                                            // 先把数据存储到 state 里面
                                            for (Map.Entry<String, Object> entry :
                                                    value.entrySet()) {
                                                if (value.containsKey("id_number")
                                                        && value.getString("id_number") != null) {
                                                    idNumberState.update(
                                                            value.getString("id_number"));
                                                } else {
                                                    loanKeyMapState.put(
                                                            value.getString("loan_key"), value);
                                                }
                                            }

                                            final String idNumber = idNumberState.value();
                                            // 直接拼接输出
                                            if (idNumber != null) {
                                                loanKeyMapState
                                                        .entries()
                                                        .forEach(
                                                                data -> {
                                                                    final JSONObject
                                                                            outPutJsonObject =
                                                                                    data.getValue();
                                                                    outPutJsonObject.put(
                                                                            "id_number", idNumber);
                                                                    if (outPutJsonObject.getString(
                                                                                            "user_key")
                                                                                    != null
                                                                            && outPutJsonObject
                                                                                            .getString(
                                                                                                    "id_number")
                                                                                    != null
                                                                            && outPutJsonObject
                                                                                            .getString(
                                                                                                    "loan_key")
                                                                                    != null) {
                                                                        // topic name 后面需要改
                                                                        outPutJsonObject.put(
                                                                                "topic_name",
                                                                                "adm_hhl01_kafka_user_repay");
                                                                        outPutJsonObject.remove(
                                                                                "stream_type");
                                                                        out.collect(
                                                                                outPutJsonObject
                                                                                        .toString(
                                                                                                SerializerFeature
                                                                                                        .WriteMapNullValue));
                                                                        try {
                                                                            // 把已经输出的数据删除掉
                                                                            loanKeyMapState.remove(
                                                                                    outPutJsonObject
                                                                                            .getString(
                                                                                                    "loan_key"));
                                                                        } catch (Exception e) {
                                                                            throw new RuntimeException(
                                                                                    e);
                                                                        }
                                                                    }
                                                                });
                                            }
                                        } catch (Exception e) {
                                            LOG.error("repayPlan 解析失败了" + e + "原始数据为: " + value);
                                        }
                                    }
                                })
                        .name("user_repay_plan");

        /** Flink Stream Sink */
        KafkaSink<String> kafkaSink =
                KafkaSink.<String>builder()
                        .setBootstrapServers(paramReader.getKfkBootStrapServer())
                        .setRecordSerializer(
                                KafkaRecordSerializationSchema.builder()
                                        .setTopicSelector(
                                                new TopicSelector<String>() {
                                                    @Override
                                                    public String apply(String record) {
                                                        return JSONObject.parseObject(record)
                                                                .getString("topic_name");
                                                    }
                                                })
                                        .setKeySerializationSchema(
                                                new SerializationSchema<String>() {
                                                    @Override
                                                    public byte[] serialize(String record) {
                                                        JSONObject jsonObject =
                                                                JSONObject.parseObject(record);
                                                        return jsonObject
                                                                .getString("user_key")
                                                                .getBytes();
                                                    }
                                                })
                                        .setValueSerializationSchema(
                                                new SerializationSchema<String>() {
                                                    @Override
                                                    public byte[] serialize(String record) {
                                                        JSONObject jsonObject =
                                                                JSONObject.parseObject(record);
                                                        jsonObject.remove("topic_name");
                                                        return jsonObject
                                                                .toString(
                                                                        SerializerFeature
                                                                                .WriteMapNullValue)
                                                                .getBytes();
                                                    }
                                                })
                                        .setPartitioner(
                                                new FlinkKafkaPartitioner<String>() {
                                                    @Override
                                                    public int partition(
                                                            String record,
                                                            byte[] key,
                                                            byte[] value,
                                                            String topic,
                                                            int[] partitions) {
                                                        return Math.abs(
                                                                new String(key).hashCode()
                                                                        % partitions.length);
                                                    }
                                                })
                                        .build())
                        .build();

        process1.union(process2).sinkTo(kafkaSink).setParallelism(1);
        env.execute(AdmUserLoanLabel.class.getSimpleName());
    }

    private static class WeKafkaKeyedSerializationSchema
            implements KeyedSerializationSchema<String> {
        @Override
        public byte[] serializeKey(String element) {
            JSONObject jsonObject = JSONObject.parseObject(element);
            String keyby = jsonObject.getString(OUTKEYBY);
            return keyby.getBytes();
        }

        @Override
        public byte[] serializeValue(String element) {
            return element.getBytes();
        }

        @Override
        public String getTargetTopic(String element) {
            return null;
        }
    }

    private static class WeKafkaCustomPartitioner extends FlinkKafkaPartitioner<String> {

        @Override
        public int partition(
                String record, byte[] key, byte[] value, String targetTopic, int[] partitions) {
            int partition = Math.abs(new String(key).hashCode() % partitions.length);
            if (LOG.isDebugEnabled()) {
                LOG.info(
                        " partitions: "
                                + partitions.length
                                + " partition: "
                                + partition
                                + " key: "
                                + new String(key));
            }
            return partition;
        }
    }

    private static JSONObject initState(String... fileds) {
        JSONObject jsonObject = new JSONObject();
        for (String filed : fileds) {
            jsonObject.put(filed, null);
        }
        return jsonObject;
    }
}
